{# ----------------------------------------------------------------------------
 # SymForce - Copyright 2022, Skydio, Inc.
 # This source code is under the Apache 2.0 license found in the LICENSE file.
 # ---------------------------------------------------------------------------- #}

#pragma once

#include <cassert>
#include <random>

#include <Eigen/Dense>

#include "../storage_ops.h"
#include {{ lcm_type_t_include_dir }}

namespace sym {

/**
 * C++ StorageOps implementation for fixed size matrices.
 */
template <typename ScalarType, int Rows, int Cols>
struct StorageOps<Eigen::Matrix<ScalarType, Rows, Cols>> {
  using Scalar = ScalarType;
  using T = Eigen::Matrix<Scalar, Rows, Cols>;
  static_assert(std::is_floating_point<ScalarType>::value, "");

  static constexpr int32_t StorageDim() {
    return Rows * Cols;
  }

  static void ToStorage(const T& a, ScalarType* out) {
    assert(out != nullptr);
    std::copy_n(a.data(), a.size(), out);
  }

  static T FromStorage(const ScalarType* data) {
    assert(data != nullptr);
    return Eigen::Map<const T>(data);
  }

  static constexpr type_t TypeEnum() {
    {% for i in range(1, 10) %}
    {% if i != 1 %}} else {% endif %}if (Rows == {{ i }} && Cols == 1) {
      return type_t::VECTOR{{ i }};
    {% endfor %}
    {% for i in range(1, 10) %}
    {% for j in range(2, 10) %}
    } else if (Rows == {{ i }} && Cols == {{ j }}) {
      return type_t::MATRIX{{ i }}{{ j }};
    {% endfor %}
    {% endfor %}
    } else {
      return type_t::INVALID;
    }
  }

  template <typename Generator>
  static T Random(Generator& gen) {
    static_assert(T::RowsAtCompileTime != Eigen::Dynamic, "Size must be known at compile time");
    static_assert(T::ColsAtCompileTime != Eigen::Dynamic, "Size must be known at compile time");

    std::normal_distribution<Scalar> distribution{};
    return T::NullaryExpr([&]() { return distribution(gen); });
  }
};

}  // namespace sym

// Explicit instantiation
{% for scalar in scalar_types %}
{% for i in range(1, 10) %}
extern template struct sym::StorageOps<Eigen::Matrix<{{ scalar }}, {{ i }}, 1>>;
{% endfor %}
{% for i in range(2, 10) %}
extern template struct sym::StorageOps<Eigen::Matrix<{{ scalar }}, {{ i }}, {{ i }}>>;
{% endfor %}
{% endfor %}
